[None][perf] Avoid implicit device-scalar syncs in DeepSeek-V4 ctx sparse metadata#16734
Conversation
|
/bot run |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
WalkthroughDeepSeek-V4 sparse attention computes and caches host-side context compressed-token totals, reuses them during eligible KV-lens updates, and retains a device-scalar fallback. A parametrized test verifies identical compressed context position IDs across both paths. ChangesDeepSeek-V4 context compressed-token sizing
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant DeepseekV4Metadata
participant prepare_compressed_kv_metadata
participant PositionIDBuilder
DeepseekV4Metadata->>prepare_compressed_kv_metadata: pass host context totals
prepare_compressed_kv_metadata->>PositionIDBuilder: forward ctx_output_sizes
PositionIDBuilder-->>prepare_compressed_kv_metadata: compute compressed position IDs
DeepseekV4Metadata->>prepare_compressed_kv_metadata: reuse cached totals for eligible KV-lens updates
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py (1)
277-277: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse complete Python 3.10+ annotations.
Use built-in generics and
| None, and add explicit return annotations to the altered production methods and added test helpers.
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py#L277-L277: usedict[int, int] | None.tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py#L745-L745: usedict[int, int] | None.tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py#L792-L797: use native annotations and add-> None.tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py#L1035-L1042: use native annotations and add-> None.tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py#L461-L463: annotate parameters and-> None.tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py#L479-L479: annotate the helper parameter and-> torch.Tensor.As per coding guidelines, “Annotate every function” and “prefer built-in generic types and
|.”🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py` at line 277, Use complete Python 3.10+ annotations throughout the listed sites: in deepseek_v4.py lines 277 and 745, replace Optional/typing generics with dict[int, int] | None; annotate the altered methods at lines 792-797 and 1035-1042 with native parameter types and -> None. In tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py lines 461-463, annotate helper parameters and return -> None, and at lines 479 annotate the helper parameter and return -> torch.Tensor.Source: Coding guidelines
tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py (1)
491-493: 🎯 Functional Correctness | 🔵 Trivial | 🏗️ Heavy liftCover the cached-size update path.
This test validates the static helper only; it never exercises
prepare()caching atdeepseek_v4.pyLines 772-774 or the reuse/fallback gate at Lines 882-885. Add a metadata-level test for bothnum_chunked_ctx_requests == 0(cached host size) and> 0(device fallback).Test coverage summary: the new test covers host/device position-ID equivalence, empty output, and compression boundaries. CI/QA test-list membership is not verifiable from the supplied files. Verdict: needs follow-up. As per path instructions, changed test functions require coverage and test-list assessment.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py` around lines 491 - 493, Extend the test near the existing golden/fast comparison to exercise the metadata flow through prepare() and the reuse/fallback gate in deepseek_v4.py. Add cases for num_chunked_ctx_requests == 0, verifying cached host-size position IDs, and greater than 0, verifying device fallback; cover host/device position-ID equivalence, empty output, and compression-boundary behavior while preserving the existing static-helper assertion.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py`:
- Around line 805-813: Update the documentation for ctx_output_sizes in
on_update_kv_lens() to state that cached host sizes are passed when no chunked
context requests exist, while the device-scalar fallback remains active only
when those cached sizes are unavailable.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py`:
- Line 277: Use complete Python 3.10+ annotations throughout the listed sites:
in deepseek_v4.py lines 277 and 745, replace Optional/typing generics with
dict[int, int] | None; annotate the altered methods at lines 792-797 and
1035-1042 with native parameter types and -> None. In
tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py
lines 461-463, annotate helper parameters and return -> None, and at lines 479
annotate the helper parameter and return -> torch.Tensor.
In
`@tests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py`:
- Around line 491-493: Extend the test near the existing golden/fast comparison
to exercise the metadata flow through prepare() and the reuse/fallback gate in
deepseek_v4.py. Add cases for num_chunked_ctx_requests == 0, verifying cached
host-size position IDs, and greater than 0, verifying device fallback; cover
host/device position-ID equivalence, empty output, and compression-boundary
behavior while preserving the existing static-helper assertion.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e40e98d2-8e38-46b6-96a4-46880c2a5759
📒 Files selected for processing (2)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py
|
PR_Github #60983 [ run ] triggered by Bot. Commit: |
|
PR_Github #60983 [ run ] completed with state
|
|
/bot run |
be52123 to
23ea630
Compare
|
PR_Github #61031 [ run ] triggered by Bot. Commit: |
|
PR_Github #61031 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61047 [ run ] triggered by Bot. Commit: |
|
PR_Github #61047 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61089 [ run ] triggered by Bot. Commit: |
|
PR_Github #61089 [ run ] completed with state
|
|
CI status note (for reviewers): pipeline #49343 is green everywhere (both builds, SBSA single-GPU) except one test — Supporting evidence that this is a main-side flake:
Re-running CI. |
|
/bot run --disable-fail-fast |
|
PR_Github #61114 [ run ] triggered by Bot. Commit: |
|
PR_Github #61114 [ run ] completed with state
|
|
/bot run |
23ea630 to
59ab788
Compare
|
PR_Github #61136 [ run ] triggered by Bot. Commit: |
|
PR_Github #61136 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61142 [ run ] triggered by Bot. Commit: |
lancelly
left a comment
lancelly
left a comment
There was a problem hiding this comment.
Overall LGTM.
Nit: The comments are a bit verbose. Maybe we can simply them.
…arse metadata DeepseekV4TrtllmAttentionMetadata._compute_ctx_compressed_position_ids sized a torch.arange and bounded a slice with a 0-dim CUDA tensor (cu_new_comp[num_contexts]). Each such implicit read is a 4-byte D2H copy plus a cudaStreamSynchronize, and with the overlap scheduler the first read per step drains the entire compute-stream queue (~250 ms of queued forward work on a GB300 disagg ctx worker). The path runs once from prepare() and twice more per step via on_update_kv_lens() (model_engine._preprocess_inputs calls it unconditionally), for 18 device-scalar reads per step in total. prepare() already computes the per-ratio ctx compressed-token counts on host tensors. Thread them through prepare_compressed_kv_metadata() into _compute_ctx_compressed_position_ids as Python ints (mirroring the gen-path gen_output_offsets precedent), and cache them on the metadata so on_update_kv_lens() can reuse them when no device-side mutation can have changed ctx-row kv_lens (num_chunked_ctx_requests == 0; only the agg extend_ctx path adjusts ctx rows on device). The device-scalar fallback is kept for that case. Validated on a DeepSeek-V4-Pro agentx disagg ctx-only GB300 benchmark (dep8, MTP3, 16384-token prefill chunks, nsys on ctx iters 210-230): whole-trace cudaStreamSynchronize count drops 360 -> 0, _prepare_inputs falls 268.3 -> 20.0 ms/step, steady-state GPU idle 282 -> 48 ms, and the executor thread's longest GIL hold shrinks 292.7 -> 19.8 ms. Position-id outputs are element-wise identical to the fallback path (unit test included; also verified sync-free under torch.cuda.set_sync_debug_mode). Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>
59ab788 to
bfe0e36
Compare
|
/bot run |
|
PR_Github #61210 [ run ] triggered by Bot. Commit: |
|
PR_Github #61142 [ run ] completed with state |
|
PR_Github #61210 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61218 [ run ] triggered by Bot. Commit: |
|
PR_Github #61218 [ run ] completed with state
|
|
Recurring CI infrastructure failure (for CI/infra visibility): this PR has now hit the same Kubernetes/JNLP agent-launch failure three times — pipelines #49288, #49388, and #49458 — all on Jenkins-Sub-3: Each occurrence kills the Build-SBSA pod before any compilation starts and fail-fast aborts the rest. The official failure analyses for all three classify it as Infrastructure/Transient with "PR likely to blame? No", and recommend escalating to CI infrastructure on recurrence — flagging it here since it has now recurred twice. No code-level blockers remain on this PR (base includes #16724 and #16761; two approvals). Re-running. |
|
/bot run |
|
PR_Github #61230 [ run ] triggered by Bot. Commit: |
|
PR_Github #61230 [ run ] completed with state
|
|
/bot run |
|
PR_Github #61256 [ run ] triggered by Bot. Commit: |
|
PR_Github #61256 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61263 [ run ] triggered by Bot. Commit: |
|
PR_Github #61263 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61295 [ run ] triggered by Bot. Commit: |
…ng on the host `request.get_tokens(0)` marshals the whole C++ VecTokens (entire sequence) into a fresh Python list of boxed ints -- O(seq_len) host work that grows linearly with ISL. Three call sites in `_prepare_tp_inputs` / `cuda_graph_runner` pay this only to read a length or a small slice, and re-pay it every iteration. In chunked prefill the context loop re-marshals the full prompt for every chunk -> O(L^2/chunk) per request over the prefill. Normal decode is already immune (get_last_tokens(0), O(1)); the waste is in the prefill/context and MTP first-draft paths. Changes: - New nanobind binding `LlmRequest.get_tokens_range(beam, begin, end)` that copies only [begin, end) (O(chunk)) instead of the whole VecTokens. - context loop and first_draft loop: use get_tokens_range for the chunk and get_num_tokens(0) (O(1)) for lengths, instead of get_tokens(0) + slice. - cuda_graph_runner first-draft branch: len(get_tokens(0)) -> get_num_tokens(0). Output is bit-identical (get_tokens_range returns the same subrange). Measured on a DSv4-Pro disagg CTX worker (c3120, non-overlap, CTX nsys): _prepare_inputs p50 27.53 ms -> 15.10 ms (-45%), p90 43.29 -> 20.43 ms; 100% request success. Complements NVIDIA#16734 (which removes the overlap device-scalar sync, 268 -> 19.96 ms); together they target the two independent costs in _prepare_inputs. Signed-off-by: Yukun He <23156053+hyukn@users.noreply.github.com> Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
PR_Github #61295 [ run ] completed with state |
…ta (upstream PR NVIDIA#16734) Squash of upstream PR NVIDIA#16734 (fixes issue NVIDIA#15684): the DSv4 sparse-attention ctx metadata prep sized a torch.arange and bounded a slice with a 0-dim CUDA tensor (cu_new_comp[num_contexts]) in _compute_ctx_compressed_position_ids, forcing an implicit D2H read + cudaStreamSynchronize. Under the overlap scheduler the first such read drains the whole compute-stream queue, which manifests as a hard deadlock (GPU 0%, iters frozen) whenever a context/prefill request is admitted into a full gen batch (MAX_UTILIZATION oversubscription); spec-algorithm-independent (MTP and DSpark both hang). Pre-extract per-ratio ctx compressed-token counts as host Python ints in prepare() and thread them (ctx_output_sizes) into prepare_compressed_kv_metadata -> _compute_ctx_compressed_position_ids so the arange size / slice bound stay off-device. on_update_kv_lens() reuses the cached host sizes EXCEPT when num_chunked_ctx_requests > 0 (the extend_ctx path may have mutated ctx-row kv_lens on device), where it falls back to the device read for correctness. Includes the upstream unit test asserting the host-int path reproduces the device-scalar fallback exactly. Supersedes the earlier local re-derivation; this is the upstream version, which adds the chunked-context staleness guard the local patch lacked. Signed-off-by: chungen28 <chung-en@deepinfra.com>
pengbowang-nv
left a comment
pengbowang-nv
left a comment
There was a problem hiding this comment.
LGTM, approve as this is targeted disagg pure ctx server.
5 participants
Background
On a DeepSeek-V4-Pro disaggregated context worker (GB300, dep8, MTP3, overlap scheduler, 16384-token prefill chunks), nsys shows
_prepare_inputstaking 268 ms/step, 92% of which is a singlecudaStreamSynchronize. The sync is triggered by implicit device-scalar reads in the DSV4 sparse-attention metadata path:_compute_ctx_compressed_position_idssizes atorch.arangeand bounds a slice with a 0-dim CUDA tensor (cu_new_comp[num_contexts]). Each such read is a 4-byte D2H copy +cudaStreamSynchronize; with the overlap scheduler the first read per step drains the entire compute-stream queue (~250 ms of queued forward work). The path executes:prepare()(2 reads x 3 compress ratios = 6 syncs), andon_update_kv_lens()—model_engine._preprocess_inputscalls it unconditionally for every batch, and the spec-decode+overlap branch calls it again — adding 12 more syncs,for 18 device-scalar reads per step. The executor thread holds the GIL for up to ~290 ms per step while blocked, delaying KV-transfer/response threads.
Changes
prepare()already computes the per-ratio ctx compressed-token counts on host tensors. Thread them throughprepare_compressed_kv_metadata()into_compute_ctx_compressed_position_idsas Python ints (ctx_output_sizes), mirroring the existing gen-pathgen_output_offsetsprecedent ("avoid tensor-scalar slice").ctx_output_sizeson the metadata soon_update_kv_lens()reuses it when no device-side mutation can have changed ctx-row kv_lens (num_chunked_ctx_requests == 0). Only the agg extend_ctx path adjusts ctx rows on device (previous_kv_lens_offsets); it falls back to the original device-scalar path.torch.cuda.set_sync_debug_mode("error")on GB300.Validation (nsys A/B on the same benchmark config, ctx iters 210-230)
cudaStreamSynchronize_prepare_inputsavgThe full agentx e2e benchmark (990K ISL conversation traces, conc 144) runs to completion with the patch (hot-patched via bind-mount before this PR); position-id outputs are bit-identical by construction and by test.
Impact assessment
num_chunked_ctx_requests > 0) and for any caller that does not passctx_output_sizes, so behavior there is unchanged.Dev Engineer Review
ctx_output_sizes) when context row lengths are known to be unchanged._ctx_output_sizes) populated duringprepare()(whennum_contexts > 0) and reused inon_update_kv_lens()only for the safe eligibility case (num_chunked_ctx_requests == 0); otherwise it preserves the existing device-scalar fallback by passingNone.prepare_compressed_kv_metadata(..., ctx_output_sizes=...)and_compute_ctx_compressed_position_ids(..., ctx_output_sizes=...)) to drive context-only compressed position-id sizing from host values while retaining the original fallback behavior for aggregate extend-context cases and callers without host counts.QA Engineer Review
test_ctx_position_ids_host_sizes_match_device_scalar_fallbackintests/unittest/_torch/attention/sparse/deepseek_v4/test_compressor_module.py.tests/integration/test_lists/,tests/qa/, ortests/waives.txtfrom the available repository inspection.